package fmi23block1;

import java.util.Collection;
import java.util.Set;

import org.jgrapht.graph.DefaultDirectedGraph;
import org.jgrapht.graph.DefaultEdge;
import org.sat4j.core.Vec;
import org.sat4j.core.VecInt;
import org.sat4j.specs.*;

/**
 * Defines the Reductions
 * 
 */
public class SatEncodings {

  private int red(int vertex) {
    return 3 * (vertex - 1) + 1;
  }

  private int green(int vertex) {
    return 3 * (vertex - 1) + 2;
  }

  private int blue(int vertex) {
    return 3 * (vertex - 1) + 3;
  }

  private int color(int literal) {
    return (literal - 1) % 3;
  }

  private int vertex(int literal) {
    return ((literal - 1) - ((literal - 1) % 3)) / 3 + 1;
  }

  /**
   * Gets a graph as input and generates a (representation of a) propositional formula that is
   * satisfiable iff the graph has a valid coloring.
   * 
   * @param graph the input graph
   * @return a cnf formula
   */
  public IVec<IVecInt> validColoring(DefaultDirectedGraph<Integer, DefaultEdge> graph) {
    // construct cnf
    IVec<IVecInt> phi = new Vec<>();

    // at least one is not red
    IVecInt phi5_dnf = new VecInt();

    for (Integer vertex : graph.vertexSet()) {
      // every vertex has exactly one color
      phi.push(new VecInt(new int[] {  red(vertex),  green(vertex),  blue(vertex) }));
      phi.push(new VecInt(new int[] { -red(vertex), -green(vertex)                }));
      phi.push(new VecInt(new int[] { -red(vertex),                 -blue(vertex) }));
      phi.push(new VecInt(new int[] {               -green(vertex), -blue(vertex) }));

      // at least one is not red
      phi5_dnf.push(-red(vertex));

      // green must support some blue vertex
      IVecInt dnf1 = new VecInt();
      dnf1.push(-green(vertex));
      Set<DefaultEdge> edges = graph.outgoingEdgesOf(vertex);
      if (!edges.isEmpty()) {
        // add green of v negatively
        for (DefaultEdge edge : edges) {
          // add blue of supported
          dnf1.push(blue(graph.getEdgeTarget(edge)));
        }
      }
      phi.push(dnf1);

      // supports more than two supports green
      if (edges.size() >= 2) {
        IVecInt dnf2 = new VecInt();
        for (DefaultEdge edge : edges) {
          // add green of supported
          dnf2.push(green(graph.getEdgeTarget(edge)));
        }
        phi.push(dnf2);
      }
    }

    // add constructed phi5 sub-dnf
    phi.push(phi5_dnf);

    for (DefaultEdge edge : graph.edgeSet()) {
      int a = graph.getEdgeSource(edge);
      int b = graph.getEdgeTarget(edge);
      // blue can only be supported by not blue
      phi.push(new VecInt(new int[]{ -blue(b), -blue(a) }));
    }

    return phi;
  }

  /**
   * Gets a graph and a set of vertices as input and generates a (representation of a) propositional
   * formula that is satisfiable iff the graph has a valid coloring that colors all the vertices in the
   * set blue.
   * 
   * @param graph the input graph
   * @param blueVertices the vertices to be colored blue
   * @return a cnf formula
   */
  public IVec<IVecInt> validColoring(DefaultDirectedGraph<Integer, DefaultEdge> graph,
      Collection<Integer> blueVertices) {
    // get original reduction phi
    IVec<IVecInt> phi = validColoring(graph);

    // add blue unit clauses for b in B
    for (Integer vertex : blueVertices) {
      phi.push(new VecInt(new int[]{ blue(vertex) }));
    }

    return phi;
  }

  /**
   * Generates an CNF that for a given instance graph and valid coloring tests whether the coloring
   * is blue-maximal. The formula is unsatisfiable iff the coloring is blue-maximal, and the models
   * of the formula correspond to valid colorings C' such that the vertices colored blue by the
   * original coloring are a strict subset of the vertices colored blue by the new coloring C'.
   * 
   * @param graph the input graph
   * @param coloring the coloring to be tested
   * @return a cnf formula
   */
  public IVec<IVecInt> maxColoring(DefaultDirectedGraph<Integer, DefaultEdge> graph,
      Coloring coloring) {
    // get original reduction phi
    IVec<IVecInt> phi = validColoring(graph);

    // generate set S of blue vertices
    IVecInt dnf = new VecInt();
    for (Integer vertex : graph.vertexSet()) {
      if (coloring.getColor(vertex) == Coloring.Color.BLUE) {
        // add unit clauses
        phi.push(new VecInt(new int[] { blue(vertex) } ));
      } else {
        // add to dnf
        dnf.push(blue(vertex));
      }
    }

    // add constructed dnf
    phi.push(dnf);

    return phi;
  }

  /**
   * Takes a model of the formula generated via one of the other methods. Returns the
   * coloring corresponding to the model.
   * 
   * @param model the model to be decoded
   * @return the corresponding coloring
   */
  public Coloring getColoringFromModel(int[] model) {
    Coloring coloring = new Coloring();
    for (int l : model) {
      boolean negated = l < 0;
      int non_negated = negated ? -l : l;
      //System.out.printf("%b %d %d %d%n", !negated, l, color(non_negated), vertex(non_negated));
      if (!negated) {
        Coloring.Color color;
        switch(color(non_negated)) {
          case 0:
            color = Coloring.Color.RED;
            break;
          case 1:
            color = Coloring.Color.GREEN;
            break;
          default:
            color = Coloring.Color.BLUE;
        }
        coloring.setColor(vertex(non_negated), color);
      }
    }
    //System.out.println(coloring);
    return coloring;
  }
}
